# model_trainer.py
import os

from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.metrics import f1_score, classification_report
import joblib
from utils import ensure_dir
from config import MODEL_SAVE_PATH


class ModelTrainer:
    def __init__(self):
        self.model = Pipeline([
            ('tfidf', TfidfVectorizer(
                max_features=5000,
                ngram_range=(1, 2),
                stop_words='english',
                sublinear_tf=True
            )),
            ('clf', LogisticRegression(
                max_iter=1000,
                class_weight='balanced',
                random_state=42,
                solver='liblinear'
            ))
        ])

    def train(self, texts, labels):
        X_train, X_val, y_train, y_val = train_test_split(
            texts, labels, test_size=0.2, random_state=42
        )

        print("开始训练模型...")
        self.model.fit(X_train, y_train)

        val_pred = self.model.predict(X_val)
        f1 = f1_score(y_val, val_pred)
        print(f"\nF1分数: {f1:.4f}")
        print("\n分类报告:")
        print(classification_report(y_val, val_pred, target_names=['人类文本(0)', 'AI生成文本(1)']))

        ensure_dir(os.path.dirname(MODEL_SAVE_PATH))
        joblib.dump(self.model, MODEL_SAVE_PATH)
        print(f"模型已保存到: {MODEL_SAVE_PATH}")